package ai.george.ai.controller.openai.rag;


import com.alibaba.cloud.ai.dashscope.chat.DashScopeChatOptions;
import io.swagger.v3.oas.annotations.Operation;
import io.swagger.v3.oas.annotations.tags.Tag;
import org.apache.commons.csv.CSVFormat;
import org.apache.commons.csv.CSVParser;
import org.apache.commons.csv.CSVRecord;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.document.Document;
import org.springframework.ai.openai.OpenAiChatModel;
import org.springframework.ai.rag.advisor.RetrievalAugmentationAdvisor;
import org.springframework.ai.rag.retrieval.search.VectorStoreDocumentRetriever;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.ai.vectorstore.filter.Filter;
import org.springframework.core.io.ClassPathResource;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;

import java.io.IOException;
import java.io.InputStreamReader;
import java.util.*;

@Tag(name = "OpenAI 智能咖啡店 RAG")
@RestController
@RequestMapping("/open-rag")
public class RAGController {


    private final VectorStore vectorStore;

    private final ChatClient chatClient;

    private static final String TYPE = "type";


    private static final String COFFEE = "coffee";

    public RAGController(VectorStore vectorStore, OpenAiChatModel chatModel) {


        this.vectorStore = vectorStore;
        Filter.Expression filterExpression = new Filter.Expression(Filter.ExpressionType.EQ, new Filter.Key(TYPE), new Filter.Value(COFFEE));

        VectorStoreDocumentRetriever vectorStoreRetriever = VectorStoreDocumentRetriever.builder()
                .vectorStore(vectorStore)
                .topK(3)
                .similarityThreshold(0.5)
                .filterExpression(filterExpression)
                .build();

        RetrievalAugmentationAdvisor advisor = RetrievalAugmentationAdvisor.builder()
                .documentRetriever(vectorStoreRetriever)
                .build();

        this.chatClient = ChatClient.builder(chatModel)
                .defaultOptions(DashScopeChatOptions.builder().withTemperature(0.7).build())
                .defaultAdvisors(advisor)
                .build();
    }


    @Operation(summary = "CSV数据导入到向量数据库")
    @GetMapping("importCSV")
    public String importCSV() {
        // 获取文件流
        ClassPathResource resource = new ClassPathResource("COFFEE-QA.csv");
        try (InputStreamReader reader = new InputStreamReader(resource.getInputStream())) {
            // 使用Apache Commons CSV解析CSV文件
            CSVParser csvParser = CSVFormat.DEFAULT
                    .builder()
                    .setHeader() // 第一行作为标题
                    .setSkipHeaderRecord(true) // 跳过标题行
                    .build()
                    .parse(reader);
            List<Document> documents = new ArrayList<>();
            // 遍历每一行记录
            for (CSVRecord record : csvParser) {
                // 获取问题和回答字段
                String question = record.get("question");
                String answer = record.get("answer");
                String category = record.get("category");
                String keywords = record.get("keywords");
                // 将问题和回答组合成文档内容
                String content = "问题: " + question + "\n回答: " + answer + "\n所属分类: " + category + "\n关键词: " + keywords;
//                String content = "问题: " + question + "\n回答: " + answer + "\n";
                Map<String, Object> metadata = new HashMap<>();
                metadata.put(TYPE, COFFEE);
                // 创建Document对象
                Document document = Document.builder()
                        .id(UUID.randomUUID().toString())
                        .text(content)
                        .metadata(metadata)
                        .build();
                documents.add(document);
            }
            // 关闭解析器
            csvParser.close();
            vectorStore.add(documents);

//            // 将文档分批存入向量数据库，每批最多10条
//            // 添加到文档列表（阿里百炼不支持超过25个的文档，具体参考源码com.alibaba.cloud.ai.dashscope.api.DashScopeApi.embeddings）
//            int batchSize = 10;
//            for (int i = 0; i < documents.size(); i += batchSize) {
//                int end = Math.min(i + batchSize, documents.size());
//                vectorStore.add(documents.subList(i, end));
//            }

            return "成功导入 " + documents.size() + " 条记录到向量数据库";

        } catch (IOException e) {
            throw new RuntimeException(e);
        }
    }


    /**
     *  rag增强回答
     * @param question
     * @return
     */
    @GetMapping("/rag-ask")
    public String ragQuery(@RequestParam(name = "question") String question) {
        // 该chatClient已经配备了RAG能力
        return chatClient
                .prompt()
                .system("你是咖啡店的服务员，你需要回答用户的问题.")
                .user(question)
                .call()
                .content();
    }



}
